import pygame
import random
import sys
from enum import Enum

# 初始化pygame
pygame.init()

# 游戏设置
SCREEN_WIDTH = 800
SCREEN_HEIGHT = 600
GRID_SIZE = 20
GRID_WIDTH = SCREEN_WIDTH // GRID_SIZE
GRID_HEIGHT = SCREEN_HEIGHT // GRID_SIZE
FPS = 10  # 初始速度

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
GREEN = (0, 255, 0)
DARK_GREEN = (0, 200, 0)
RED = (255, 0, 0)
BLUE = (0, 0, 255)
YELLOW = (255, 255, 0)
PURPLE = (128, 0, 128)
ORANGE = (255, 165, 0)
GRAY = (128, 128, 128)

# 方向定义
class Direction(Enum):
    UP = (0, -1)
    DOWN = (0, 1)
    LEFT = (-1, 0)
    RIGHT = (1, 0)

# 食物类型
class FoodType(Enum):
    NORMAL = 1
    GOLDEN = 2
    SPEED_UP = 3
    SLOW_DOWN = 4
    DOUBLE = 5

class Snake:
    def __init__(self):
        self.reset()

    def reset(self):
        # 初始化蛇在屏幕中央
        start_x = GRID_WIDTH // 2
        start_y = GRID_HEIGHT // 2
        self.body = [(start_x, start_y), (start_x - 1, start_y), (start_x - 2, start_y)]
        self.direction = Direction.RIGHT
        self.grow_count = 0
        self.speed = FPS
        self.score = 0
        self.high_score = 0
        self.game_over = False
        self.game_started = False
        self.super_mode = 0  # 无敌模式时间

    def move(self):
        if self.game_over or not self.game_started:
            return

        # 获取头部位置
        head = self.body[0]
        dx, dy = self.direction.value
        new_head = (head[0] + dx, head[1] + dy)

        # 检查边界碰撞
        if (new_head[0] < 0 or new_head[0] >= GRID_WIDTH or
            new_head[1] < 0 or new_head[1] >= GRID_HEIGHT):
            if self.super_mode <= 0:  # 无敌模式下可以穿墙
                self.game_over = True
                return
            else:
                # 无敌模式下穿墙到对面
                new_head = (new_head[0] % GRID_WIDTH, new_head[1] % GRID_HEIGHT)

        # 检查自身碰撞
        if new_head in self.body[1:] and self.super_mode <= 0:
            self.game_over = True
            return

        # 移动蛇
        self.body.insert(0, new_head)

        # 如果需要生长，不移除尾部
        if self.grow_count > 0:
            self.grow_count -= 1
        else:
            self.body.pop()

        # 更新无敌模式
        if self.super_mode > 0:
            self.super_mode -= 1

    def change_direction(self, new_direction):
        # 防止蛇往相反方向移动
        if self.game_over or not self.game_started:
            return

        opposite = {
            Direction.UP: Direction.DOWN,
            Direction.DOWN: Direction.UP,
            Direction.LEFT: Direction.RIGHT,
            Direction.RIGHT: Direction.LEFT
        }

        if new_direction != opposite.get(self.direction):
            if new_direction != self.direction:  # 只有当新方向与当前方向不同时才处理
                self.direction = new_direction

    def grow(self, amount=1):
        self.grow_count += amount

    def check_food_collision(self, food):
        return self.body[0] == food.position

    def draw(self, screen):
        for i, segment in enumerate(self.body):
            # 头部用不同颜色
            if i == 0:
                color = DARK_GREEN
                if self.super_mode > 0 and self.super_mode % 10 < 5:  # 无敌模式闪烁
                    color = YELLOW
            else:
                # 身体渐变颜色
                green_value = max(100, 255 - i * 10)
                color = (0, green_value, 0)

            rect = pygame.Rect(segment[0] * GRID_SIZE, segment[1] * GRID_SIZE,
                             GRID_SIZE - 2, GRID_SIZE - 2)
            pygame.draw.rect(screen, color, rect)

            # 绘制蛇眼（仅头部）
            if i == 0:
                eye_size = 3
                eye_offset = 4
                if self.direction == Direction.RIGHT:
                    eye1_pos = (segment[0] * GRID_SIZE + GRID_SIZE - eye_offset,
                              segment[1] * GRID_SIZE + eye_offset)
                    eye2_pos = (segment[0] * GRID_SIZE + GRID_SIZE - eye_offset,
                              segment[1] * GRID_SIZE + GRID_SIZE - eye_offset)
                elif self.direction == Direction.LEFT:
                    eye1_pos = (segment[0] * GRID_SIZE + eye_offset,
                              segment[1] * GRID_SIZE + eye_offset)
                    eye2_pos = (segment[0] * GRID_SIZE + eye_offset,
                              segment[1] * GRID_SIZE + GRID_SIZE - eye_offset)
                elif self.direction == Direction.UP:
                    eye1_pos = (segment[0] * GRID_SIZE + eye_offset,
                              segment[1] * GRID_SIZE + eye_offset)
                    eye2_pos = (segment[0] * GRID_SIZE + GRID_SIZE - eye_offset,
                              segment[1] * GRID_SIZE + eye_offset)
                else:  # DOWN
                    eye1_pos = (segment[0] * GRID_SIZE + eye_offset,
                              segment[1] * GRID_SIZE + GRID_SIZE - eye_offset)
                    eye2_pos = (segment[0] * GRID_SIZE + GRID_SIZE - eye_offset,
                              segment[1] * GRID_SIZE + GRID_SIZE - eye_offset)

                pygame.draw.circle(screen, WHITE, eye1_pos, eye_size)
                pygame.draw.circle(screen, WHITE, eye2_pos, eye_size)
                pygame.draw.circle(screen, BLACK, eye1_pos, eye_size - 1)
                pygame.draw.circle(screen, BLACK, eye2_pos, eye_size - 1)

class Food:
    def __init__(self, food_type=FoodType.NORMAL):
        self.type = food_type
        self.position = self.generate_position()
        self.animation_timer = 0

    def generate_position(self):
        # 生成不在蛇身上的随机位置
        while True:
            x = random.randint(0, GRID_WIDTH - 1)
            y = random.randint(0, GRID_HEIGHT - 1)
            return (x, y)

    def draw(self, screen):
        self.animation_timer += 1

        # 根据食物类型选择颜色和形状
        if self.type == FoodType.NORMAL:
            color = RED
            size = GRID_SIZE - 4
            # 普通食物 - 圆形
            center = (self.position[0] * GRID_SIZE + GRID_SIZE // 2,
                     self.position[1] * GRID_SIZE + GRID_SIZE // 2)
            pygame.draw.circle(screen, color, center, size // 2)

        elif self.type == FoodType.GOLDEN:
            color = YELLOW
            size = GRID_SIZE - 2
            # 金色食物 - 星形
            center = (self.position[0] * GRID_SIZE + GRID_SIZE // 2,
                     self.position[1] * GRID_SIZE + GRID_SIZE // 2)
            self.draw_star(screen, center, size // 2, color)

        elif self.type == FoodType.SPEED_UP:
            color = BLUE
            size = GRID_SIZE - 4
            # 加速食物 - 三角形向上
            points = [
                (self.position[0] * GRID_SIZE + GRID_SIZE // 2, self.position[1] * GRID_SIZE + 2),
                (self.position[0] * GRID_SIZE + 2, self.position[1] * GRID_SIZE + GRID_SIZE - 2),
                (self.position[0] * GRID_SIZE + GRID_SIZE - 2, self.position[1] * GRID_SIZE + GRID_SIZE - 2)
            ]
            pygame.draw.polygon(screen, color, points)

        elif self.type == FoodType.SLOW_DOWN:
            color = PURPLE
            size = GRID_SIZE - 4
            # 减速食物 - 三角形向下
            points = [
                (self.position[0] * GRID_SIZE + GRID_SIZE // 2, self.position[1] * GRID_SIZE + GRID_SIZE - 2),
                (self.position[0] * GRID_SIZE + 2, self.position[1] * GRID_SIZE + 2),
                (self.position[0] * GRID_SIZE + GRID_SIZE - 2, self.position[1] * GRID_SIZE + 2)
            ]
            pygame.draw.polygon(screen, color, points)

        elif self.type == FoodType.DOUBLE:
            color = ORANGE
            size = GRID_SIZE - 4
            # 双倍分食物 - 菱形
            center = (self.position[0] * GRID_SIZE + GRID_SIZE // 2,
                     self.position[1] * GRID_SIZE + GRID_SIZE // 2)
            points = [
                (center[0], center[1] - size // 2),
                (center[0] + size // 2, center[1]),
                (center[0], center[1] + size // 2),
                (center[0] - size // 2, center[1])
            ]
            pygame.draw.polygon(screen, color, points)

        # 添加脉冲效果
        if self.type in [FoodType.GOLDEN, FoodType.DOUBLE]:
            pulse = abs(math.sin(self.animation_timer * 0.1)) * 3
            pygame.draw.rect(screen, color,
                           (self.position[0] * GRID_SIZE - pulse,
                            self.position[1] * GRID_SIZE - pulse,
                            GRID_SIZE + pulse * 2, GRID_SIZE + pulse * 2), 1)

    def draw_star(self, screen, center, size, color):
        # 绘制五角星
        points = []
        for i in range(10):
            angle = math.pi * i / 5
            if i % 2 == 0:
                radius = size
            else:
                radius = size // 2
            x = center[0] + radius * math.cos(angle - math.pi / 2)
            y = center[1] + radius * math.sin(angle - math.pi / 2)
            points.append((x, y))
        pygame.draw.polygon(screen, color, points)

class Game:
    def __init__(self):
        self.snake = Snake()
        self.foods = []
        self.font_large = pygame.font.Font(None, 48)
        self.font_medium = pygame.font.Font(None, 36)
        self.font_small = pygame.font.Font(None, 24)
        self.clock = pygame.time.Clock()
        self.screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
        pygame.display.set_caption("贪吃蛇 - Snake Game")

        self.spawn_food()

    def spawn_food(self):
        # 清除已存在的食物
        self.foods.clear()

        # 根据分数生成不同类型的食物
        food_type = FoodType.NORMAL

        # 随机生成特殊食物
        rand = random.random()
        if self.snake.score >= 50 and rand < 0.1:  # 金色食物
            food_type = FoodType.GOLDEN
        elif self.snake.score >= 30 and rand < 0.15:  # 双倍分食物
            food_type = FoodType.DOUBLE
        elif self.snake.score >= 20 and rand < 0.2:  # 速度变化食物
            food_type = random.choice([FoodType.SPEED_UP, FoodType.SLOW_DOWN])

        new_food = Food(food_type)

        # 确保食物不在蛇身上
        while new_food.position in self.snake.body:
            new_food = Food(food_type)

        self.foods.append(new_food)

    def handle_food_collision(self):
        for food in self.foods[:]:
            if self.snake.check_food_collision(food):
                # 根据食物类型给予不同的奖励
                if food.type == FoodType.NORMAL:
                    self.snake.score += 10
                    self.snake.grow(1)
                elif food.type == FoodType.GOLDEN:
                    self.snake.score += 50
                    self.snake.grow(3)
                    self.snake.super_mode = 50  # 5秒无敌时间
                elif food.type == FoodType.DOUBLE:
                    self.snake.score += 20
                    self.snake.grow(2)
                elif food.type == FoodType.SPEED_UP:
                    self.snake.score += 15
                    self.snake.grow(1)
                    self.snake.speed = max(5, self.snake.speed - 2)  # 加速
                elif food.type == FoodType.SLOW_DOWN:
                    self.snake.score += 15
                    self.snake.grow(1)
                    self.snake.speed = min(20, self.snake.speed + 3)  # 减速

                # 更新最高分
                if self.snake.score > self.snake.high_score:
                    self.snake.high_score = self.snake.score

                # 生成新食物
                self.spawn_food()

                # 随机生成额外食物（高分时）
                if self.snake.score >= 100 and random.random() < 0.3:
                    self.spawn_food()

                return True
        return False

    def draw_grid(self):
        # 绘制网格线
        for x in range(0, SCREEN_WIDTH, GRID_SIZE):
            pygame.draw.line(self.screen, GRAY, (x, 0), (x, SCREEN_HEIGHT), 1)
        for y in range(0, SCREEN_HEIGHT, GRID_SIZE):
            pygame.draw.line(self.screen, GRAY, (0, y), (SCREEN_WIDTH, y), 1)

    def draw_start_screen(self):
        self.screen.fill(BLACK)

        # 标题
        title = self.font_large.render("SNAKE GAME", True, GREEN)
        title_rect = title.get_rect(center=(SCREEN_WIDTH // 2, 150))
        self.screen.blit(title, title_rect)

        # 说明
        instructions = [
            "Use ARROW KEYS to move",
            "Eat food to grow longer",
            "Avoid walls and yourself!",
            "",
            "Food Types:",
            "RED: Normal (+10, +1 length)",
            "YELLOW: Golden (+50, +3 length, 5s invincible)",
            "ORANGE: Double (+20, +2 length)",
            "BLUE: Speed Up (+15, faster)",
            "PURPLE: Slow Down (+15, slower)",
            "",
            "Press SPACE to Start",
            "Press ESC to Quit"
        ]

        y = 220
        for instruction in instructions:
            color = WHITE
            if "RED:" in instruction:
                color = RED
            elif "YELLOW:" in instruction:
                color = YELLOW
            elif "ORANGE:" in instruction:
                color = ORANGE
            elif "BLUE:" in instruction:
                color = BLUE
            elif "PURPLE:" in instruction:
                color = PURPLE

            text = self.font_small.render(instruction, True, color)
            text_rect = text.get_rect(center=(SCREEN_WIDTH // 2, y))
            self.screen.blit(text, text_rect)
            y += 25

    def draw_game_over_screen(self):
        # 半透明覆盖
        overlay = pygame.Surface((SCREEN_WIDTH, SCREEN_HEIGHT))
        overlay.set_alpha(200)
        overlay.fill(BLACK)
        self.screen.blit(overlay, (0, 0))

        # 游戏结束文字
        game_over_text = self.font_large.render("GAME OVER", True, RED)
        game_over_rect = game_over_text.get_rect(center=(SCREEN_WIDTH // 2, 200))
        self.screen.blit(game_over_text, game_over_rect)

        # 分数显示
        score_text = self.font_medium.render(f"Score: {self.snake.score}", True, WHITE)
        score_rect = score_text.get_rect(center=(SCREEN_WIDTH // 2, 260))
        self.screen.blit(score_text, score_rect)

        high_score_text = self.font_medium.render(f"High Score: {self.snake.high_score}", True, YELLOW)
        high_score_rect = high_score_text.get_rect(center=(SCREEN_WIDTH // 2, 300))
        self.screen.blit(high_score_text, high_score_rect)

        # 重新开始提示
        restart_text = self.font_small.render("Press SPACE to Restart", True, WHITE)
        restart_rect = restart_text.get_rect(center=(SCREEN_WIDTH // 2, 360))
        self.screen.blit(restart_text, restart_rect)

        quit_text = self.font_small.render("Press ESC to Quit", True, WHITE)
        quit_rect = quit_text.get_rect(center=(SCREEN_WIDTH // 2, 390))
        self.screen.blit(quit_text, quit_rect)

    def draw(self):
        self.screen.fill(BLACK)

        if not self.snake.game_started:
            self.draw_start_screen()
        else:
            # 绘制网格
            self.draw_grid()

            # 绘制食物
            for food in self.foods:
                food.draw(self.screen)

            # 绘制蛇
            self.snake.draw(self.screen)

            # 显示分数
            score_text = self.font_medium.render(f"Score: {self.snake.score}", True, WHITE)
            self.screen.blit(score_text, (10, 10))

            high_score_text = self.font_small.render(f"High: {self.snake.high_score}", True, YELLOW)
            self.screen.blit(high_score_text, (10, 50))

            # 显示长度
            length_text = self.font_small.render(f"Length: {len(self.snake.body)}", True, GREEN)
            self.screen.blit(length_text, (10, 80))

            # 显示速度
            speed_text = self.font_small.render(f"Speed: {21 - self.snake.speed}", True, WHITE)
            self.screen.blit(speed_text, (10, 110))

            # 无敌模式提示
            if self.snake.super_mode > 0:
                invincible_text = self.font_small.render(f"INVINCIBLE: {self.snake.super_mode // 10}s", True, YELLOW)
                self.screen.blit(invincible_text, (SCREEN_WIDTH // 2 - 70, 10))

            if self.snake.game_over:
                self.draw_game_over_screen()

    def handle_events(self):
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return False

            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    return False

                if not self.snake.game_started:
                    if event.key == pygame.K_SPACE:
                        self.snake.game_started = True
                elif self.snake.game_over:
                    if event.key == pygame.K_SPACE:
                        self.snake.reset()
                        self.foods.clear()
                        self.spawn_food()
                else:
                    # 立即响应方向键按下
                    if event.key == pygame.K_UP:
                        self.snake.change_direction(Direction.UP)
                    elif event.key == pygame.K_DOWN:
                        self.snake.change_direction(Direction.DOWN)
                    elif event.key == pygame.K_LEFT:
                        self.snake.change_direction(Direction.LEFT)
                    elif event.key == pygame.K_RIGHT:
                        self.snake.change_direction(Direction.RIGHT)

        return True

    def run(self):
        running = True
        while running:
            running = self.handle_events()

            if self.snake.game_started and not self.snake.game_over:
                self.snake.move()
                self.handle_food_collision()

            self.draw()
            pygame.display.flip()
            self.clock.tick(self.snake.speed)

        pygame.quit()
        sys.exit()

# 导入数学模块用于星形绘制
import math

# 运行游戏
if __name__ == "__main__":
    game = Game()
    game.run()